import numpy as np
from load_data import read_data, gen_dict


def viterbi(A, B, pi, O):
    '''
    :param A: 状态转移概率矩阵：一个状态转移到另一个状态的概率
    :param B: 观测概率矩阵：行表示某个状态，列表示某个观测，Bij表示在这个状态下产生这个观测的概率,
    设第一列观测序号为0，第二列观测的序号为1，...
    :param pi: 初始状态概率向量：以某个状态开始的概率,pi的长度即为可选的状态的个数
    :param O: 观测序列，由0，1，2...组成，O的长度即为观测次数
    :return: 最优路径I:一个数组，表示每个时刻最可能选择的状态的编号i，为0,1,2...
    '''

    #第一步：初始化:变量
    # k表示观测次数
    k = O.shape[0]
    # n表示可选状态的数目
    n = pi.shape[0]
    # node记录时刻t时，状态为i的所有路径中概率最大的路径的上一个结点，即为上一个状态的编号0，1，2...，大小为 k * n,每个元素初始化为0
    node = [[0 for i in range(n)] for j in range(k)]
    #delta记录时刻t时，状态i观测到O[t]的概率，大小为k * n，每个元素初始化为0
    delta = [[0 for i in range(n)] for j in range(k)]

    #第二部：单独初始化时刻1，时刻1的观测为O[0]，
    for i in range(n):
        delta[0][i] = pi[i] * B[i][O[0]]

    #第三步：递推计算:时刻2，3，4，....
    for t in range(1, k):
        for i in range(n):
            # 计算t时刻，状态i，在上一个时刻各个状态下得到的观测的基础上，该时刻i状态得到新观测值的最大概率
            # 即计算delta[t][i]
            # 计算node[t][i]，记录的是前一个状态j
            prob_max = 0
            j_max = 0
            for j in range(n):
                p = delta[t-1][j] * A[j][i]
                if p > prob_max:
                    prob_max = p
                    j_max = j
            delta[t][i] = prob_max * B[i][O[t]]
            node[t][i] = j_max

    #第四步：找到最优路径的终点
    #I为返回结果，是最优路径序列
    I = []
    i = np.argmax(delta[k-1])
    I.append(i)

    #第五步：由最优路径的终点回溯找到其他结点
    for t in range(k-2, -1, -1):
        I.insert(0, node[t+1][i])
        i = node[t+1][i]

    return np.array(I)


def ner_train(train_data, tag_to_ids, dict):

    #标注作为可选状态，共有9个状态
    #英文单词作为观测，观测个数为字典的长度
    #pi初始状态概率向量：即训练数据集中各个标签的占比
    #A状态转移矩阵：训练数据集中各个标签和下一个紧挨着的标签之间的转移数目占总转移数目的比值

    '''
    读取数据，初始化A,B,pi,
    '''
    epsilon = 1e-5
    sentence_word_ids_list, sentence_tags_list = read_data(train_data, tag_to_ids, dict)
    # n可选状态数目
    n = len(tag_to_ids)
    # d可观测值数目，也就是字典的长度
    d = len(dict)
    # 定义A,B,pi
    pi = np.zeros(n)
    A = np.zeros((n, n))
    B = np.zeros((n, d))

    for i in range(len(sentence_tags_list)):
        #生成状态转移矩阵：状态转移有：0-1 1-2 2-3 ,..., n-2-n-1
        tags_length = len(sentence_tags_list[i])
        if tags_length <= 1:
            continue

        pi[sentence_tags_list[i][0]] += 1

        for j in range(tags_length-1):
            A[sentence_tags_list[i][j]][sentence_tags_list[i][j+1]] += 1
            B[sentence_tags_list[i][j]][sentence_word_ids_list[i][j]] += 1
        B[sentence_tags_list[i][tags_length-1]][sentence_word_ids_list[i][tags_length-1]] += 1

    #把pi转化为概率
    pi[pi == 0] = epsilon
    pi /= np.sum(pi)

    #把A转化为概率
    A[A == 0] = epsilon
    A /= np.sum(A, axis=1, keepdims=True)

    #把B转化为概率
    B[B == 0] = epsilon
    B /= np.sum(B, axis=1, keepdims=True)
    return A, B, pi

def ner_test(A, B, pi, test_data, tag_to_ids, dict):

    sentence_word_ids_list, sentence_tags_list = read_data(test_data, tag_to_ids, dict)
    sentence_num = len(sentence_tags_list)
    tp = 0
    fn = 0
    fp = 0
    for i in range(sentence_num):
        if len(sentence_tags_list[i]) <= 1:
            continue
        test_res = viterbi(A, B, pi, np.array(sentence_word_ids_list[i]))
        #tp += np.sum(test_res == np.array(sentence_tags_list[i]))
        for j in range(test_res.shape[0]):
            # 计算本来是命名实体，但是被识别为O的个数
            if test_res[j] == tag_to_ids.get("O") and sentence_tags_list[i][j] != tag_to_ids.get("O"):
                fn += 1
            if test_res[j] != tag_to_ids.get("O") and sentence_tags_list[i][j] != tag_to_ids.get("O") and test_res[j] != sentence_tags_list[i][j]:
                fp += 1
            if test_res[j] == sentence_tags_list[i][j] and test_res[j] != tag_to_ids.get("O"):
                tp += 1
    p = tp / (tp + fp)
    r = tp / (tp + fn)
    f1 = 2 * p * r / (p + r)
    print("Precision为: {}".format(p))
    print("Recall为: {}".format(r))
    print("F1为: {}".format(f1))
if __name__ == "__main__":
    train_data = "./data/train"
    testa_data = "./data/testa"
    testb_data = "./data/testb"
    tag_to_ids = {"B-ORG": 0, "I-ORG": 1, "B-PER": 2, "I-PER": 3, "B-LOC": 4, "I-LOC": 5, "B-MISC": 6, "I-MISC": 7, "O": 8}

    dict = gen_dict([train_data, testa_data, testb_data])

    A, B, pi = ner_train([train_data], tag_to_ids, dict)
    ner_test(A, B, pi, [testa_data], tag_to_ids, dict)



